#! /usr/bin/env python3
#coding=utf-8
from sensor_msgs.msg import LaserScan
from mcq_msgs.msg import Carry

import signal
import rospy
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
import matplotlib.patches as patches
from matplotlib.animation import FuncAnimation, FFMpegWriter
import matplotlib.animation as animation
import numpy as np
import threading
import rospy
from sensor_msgs.msg import LaserScan
import time
angles = []
distances = []  
cluster = []
centerX = []
centerY = []  
radius = []
def DB_callback(msg):
    global angles,distances,cluster,centerX,centerY,radius
    angles = list(msg.angles)
    distances = list(msg.DB_data)
    cluster = list(msg.cluster)
    # for i in range(len(angles)):
    #     print(angles[i],distances[i],cluster[i])

def spin():
    rospy.spin()

def plot_show():
    ani = animation.FuncAnimation(fig, update, frames=range(10000), interval=10)
    ani.save('/home/mcq/ros1_lidar/src/rplidar_ros/scripts/video/scatter_animation6.mp4', writer='ffmpeg', fps=10)
    sub_DBdata.unregister()  # 取消订阅话题
    plt.close(fig)   
    
def update(frame):
    if angles and distances and cluster:
        ax.clear()
        # 使用颜色映射
        norm = mcolors.Normalize(vmin=min(cluster), vmax=max(cluster))
        cmap = plt.cm.get_cmap('viridis')
        colors = [cmap(norm(category)) if category != 0 else 'red' for category in cluster]
        # 绘制直角坐标系下的散点图
        ax.scatter(angles, distances, c=colors)
        # 设置坐标轴范围，以 (0, 0) 为中心
        ax.set_xlim(-maximum_distance, maximum_distance)
        ax.set_ylim(-maximum_distance, maximum_distance)
        #  设置 x 和 y 轴的刻度间距
        x_ticks = np.arange(-maximum_distance, maximum_distance + 1, step=50)
        y_ticks = np.arange(-maximum_distance, maximum_distance + 1, step=50)
        ax.set_xticks(x_ticks)
        ax.set_yticks(y_ticks)
        # 添加文本标注
        for category in set(cluster):
            # 获取该类别的索引
            indices = [i for i, value in enumerate(cluster) if value == category]
            if(category == 0):continue
            # 计算该类别的平均位置
            avg_x = np.mean([angles[idx] for idx in indices])
            avg_y = np.mean([distances[idx] for idx in indices])
            # 添加标签
            ax.text(avg_x, avg_y, f' {category}', fontsize=12, color='black', ha='center', va='center')
            # first_point_index = next((i for i, value in enumerate(cluster) if value == category), None)
            # if first_point_index is not None:
            #     # 获取第一个点的坐标
            #     first_point_x = angles[first_point_index]
            #     first_point_y = distances[first_point_index]
            #     # 添加标签
            #     ax.text(first_point_x, first_point_y, f' {category}', fontsize=24, color='black', ha='center', va='center')
        # 添加网格线
        ax.grid(True)
        ax.set_xlabel('X轴')
        ax.set_ylabel('Y轴')
        # ax.legend() 
        plt.pause(0.001)


if __name__ == '__main__':
    # try:
        rospy.init_node("take_DB_node",anonymous=True)
        sub_DBdata = rospy.Subscriber("/mcq_DB_topic",Carry,DB_callback,queue_size=4000)

        maximum_distance = 400
        fig, ax = plt.subplots(figsize=(10, 10))
        ax.set_aspect('equal',adjustable='datalim')

        ros_thread = threading.Thread(target=spin)
        plot_thread = threading.Thread(target=plot_show)
        
        ros_thread.start()
        plot_thread.start()
    # except BrokenPipeError:
    #     # This block will be executed when Ctrl+C is pressed
    #     rospy.loginfo("断开")
    #     sub_DBdata.unregister()

